#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import base64
from io import BytesIO
from flask import Flask, request, jsonify, send_file
import numpy as np
import cv2
from paddleocr import PaddleOCR, draw_ocr
from PIL import Image
import uuid
from datetime import datetime

def preprocess_image_for_ocr(image):
    return image
    """图像预处理，提高OCR识别率"""
    # 转为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
    
    # 自适应阈值处理
    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                 cv2.THRESH_BINARY, 15, 8)
    
    # 降噪和锐化
    kernel = np.ones((1, 1), np.uint8)
    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
    kernel_sharpen = np.array([[-1,-1,-1],
                              [-1, 9,-1],
                              [-1,-1,-1]])
    sharpen = cv2.filter2D(opening, -1, kernel_sharpen)
    
    # 增加对比度
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    return clahe.apply(sharpen)

def enhance_handwriting(image):
    return image
    # return image
    """手写体图像预处理增强"""
    # 转灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
    
    kernel_sharpen = np.array([[-1,-1,-1],
                              [-1, 9,-1],
                              [-1,-1,-1]])
    sharpen = cv2.filter2D(gray, -1, kernel_sharpen)
    
    # 增强对比度
    clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8,8))
    return clahe.apply(sharpen)

# 初始化Flask应用
app = Flask(__name__)
app.config['JSON_AS_ASCII'] = False

# CORS支持
@app.after_request
def after_request(response):
    response.headers.update({
        'Access-Control-Allow-Origin': '*',
        'Access-Control-Allow-Methods': 'GET,POST',
        'Access-Control-Allow-Headers': 'Content-Type'
    })
    return response

# 创建临时文件目录
TEMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'temp')
os.makedirs(TEMP_DIR, exist_ok=True)

# 设置字体路径
FONT_PATHS = [
    '/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc',
    '/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc',
    '/usr/share/fonts/noto-cjk/NotoSansCJK-Regular.ttc',
    '/workspace/PaddleOCR/doc/fonts/simfang.ttf',
    '/workspace/simfang.ttf'
]

FONT_PATH = next((path for path in FONT_PATHS if os.path.exists(path)), None)
if not FONT_PATH:
    print("警告: 未找到合适的字体文件，使用paddleocr默认字体")

def cleanup_old_files():
    """清理temp目录下非当日的文件"""
    try:
        current_date = datetime.now().strftime('%Y%m%d')
        for filename in os.listdir(TEMP_DIR):
            if filename.startswith('ocr_result_'):
                file_date = filename.split('_')[2]
                if file_date != current_date:
                    try:
                        os.remove(os.path.join(TEMP_DIR, filename))
                        print(f"已删除过期文件: {filename}")
                    except Exception as e:
                        print(f"删除文件失败 {filename}: {str(e)}")
    except Exception as e:
        print(f"清理过期文件时出错: {str(e)}")

def get_ocr(lang="ch"):
    """获取OCR实例"""
    return PaddleOCR(
        use_angle_cls=True,
        use_gpu=False,
        lang=lang,
        show_log=False,
        use_dilation=True,
        det_db_box_thresh=0.3,
        det_db_thresh=0.2,
        det_db_unclip_ratio=1.8,
        rec_batch_num=10,
        max_text_length=50,
        det_limit_side_len=2880,
        det_limit_type='max',
        rec_image_shape="3, 48, 320"
    )

def get_handwriting_ocr(lang="ch"):
    """获取手写体OCR实例"""
    model_config = {
        "ch": {
            "det_model_dir": '/workspace/models/ch_PP-OCRv4_det_infer',
            "rec_model_dir": '/workspace/models/ch_PP-OCRv4_rec_server_infer',
            # "rec_char_dict_path": '/workspace/dict/ppocr_keys_v1.txt',
            "cls_model_dir": '/workspace/models/ch_ppocr_mobile_v2.0_cls_infer'
        },
        "en": {
            "det_model_dir": '/workspace/models/en_PP-OCRv3_det_infer',
            "rec_model_dir": '/workspace/models/en_PP-OCRv4_rec_infer',
            "rec_char_dict_path": '/workspace/dict/en_dict.txt',
            "cls_model_dir": '/workspace/models/ch_ppocr_mobile_v2.0_cls_infer'  # 中英文共用方向分类模型
        }
    }
    
    return PaddleOCR(
        use_angle_cls=True,  # 使用方向分类
        lang=lang,
        use_gpu=False,
        show_log=False,
        **model_config[lang],
        rec_image_shape="3, 48, 320",  # 提高识别图像尺寸
        det_db_box_thresh=0.3,  # 提高检测框阈值，减少误检
        det_db_thresh=0.2,  # 增加DB检测阈值
        det_db_unclip_ratio=1.5,  # 增大文本框扩张比例
        rec_batch_num=10,  # 增加识别批处理大小
        max_text_length=80,  # 增加最大文本长度
        det_limit_side_len=4096,  # 保持大尺寸图像的细节
        det_limit_type='max',
        min_box_size=5,  # 增加最小检测框尺寸，过滤小噪点
        drop_score=0.5,  # 提高识别置信度阈值
        det_db_score_mode="slow",  # 使用更精确的评分模式
        det_db_box_type="quad",
        cls_batch_num=8,  # 增加方向分类批处理大小
        cls_thresh=0.8,  # 提高方向分类置信度阈值
        vis_font_path=FONT_PATH,  # 指定可视化字体
        rec_algorithm="SVTR_LCNet",  # 指定更高精度的识别算法
        use_dilation=True,  # 使用膨胀操作增强文本区域
        rec_char_type=lang,  # 指定字符类型
        use_space_char=True  # 使用空格字符
    )

def process_ocr_result(output_type, result_list, image, boxes=None, txts=None, scores=None):
    """处理OCR结果"""
    # 处理空结果
    if not result_list:
        result_dict = {
            "results": [],
            "message": "未检测到任何文本"
        }
        if output_type == "base64":
            buffer = BytesIO()
            pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            pil_image.save(buffer, format="JPEG")
            result_dict["image"] = base64.b64encode(buffer.getvalue()).decode('utf-8')
        return result_dict

    # 如果不需要图像输出
    if output_type == "none":
        return {"results": result_list}

    # 处理图像结果
    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    if boxes and txts and scores:
        draw_img = draw_ocr(
            image=pil_image,
            boxes=boxes,
            txts=txts,
            scores=scores,
            font_path=FONT_PATH
        )
        draw_pil = Image.fromarray(draw_img)
    else:
        draw_pil = pil_image

    # 返回base64编码的图像
    if output_type == "base64":
        buffer = BytesIO()
        draw_pil.save(buffer, format="JPEG")
        return {
            "image": base64.b64encode(buffer.getvalue()).decode('utf-8'),
            "results": result_list
        }

    # 保存为文件并返回文件路径
    current_date = datetime.now().strftime('%Y%m%d')
    filename = f"ocr_result_{current_date}_{str(uuid.uuid4()).replace('-', '')[:8]}.jpg"
    temp_file = os.path.join(TEMP_DIR, filename)
    draw_pil.save(temp_file)
    return {
        "file_path": f"/temp/{filename}",
        "results": result_list,
        "message": "注意：1.文件链接仅可访问一次，访问后文件将被自动删除；2.为确保安全，访问时会自动清理非当日的所有文件。"
    }

def validate_request():
    """验证请求数据"""
    if not request.is_json:
        return None, "请求头必须是application/json"
    
    data = request.json
    if not data:
        return None, "请求体不能为空"
    
    if "images" not in data:
        return None, "缺少必需的images字段"
    
    # 验证语言参数
    lang = data.get("lang", "ch")
    if lang not in ["ch", "en"]:
        return None, "不支持的语言类型，仅支持'ch'或'en'"
    
    # 验证输出类型
    output_type = data.get("output_type", "base64")
    if output_type not in ["base64", "file", "none"]:
        return None, f"output_type参数值'{output_type}'无效，必须是'base64'、'file'或'none'"
    
    return data, None

def process_image(image_base64, ocr_instance, output_type, preprocess_func=None, scale_ratio=None):
    """处理单个图像
    
    Args:
        image_base64: base64编码的图像数据
        ocr_instance: OCR实例
        output_type: 输出类型
        preprocess_func: 预处理函数
        scale_ratio: 缩放比例，如果提供，将对图像进行缩放处理
    """
    if not isinstance(image_base64, str):
        raise ValueError("images数组中的元素必须是base64字符串")

    # 解码图像
    image_data = base64.b64decode(image_base64)
    image = cv2.imdecode(np.frombuffer(image_data, np.uint8), cv2.IMREAD_COLOR)
    
    # 记录原始尺寸
    orig_h, orig_w = image.shape[:2]
    
    # 如果需要缩放
    if scale_ratio:
        # 计算新尺寸
        aspect_ratio = orig_w / orig_h
        if aspect_ratio > 1:
            new_w = int(orig_w * scale_ratio)
            new_h = int(new_w / aspect_ratio)
        else:
            new_h = int(orig_h * scale_ratio)
            new_w = int(new_h * aspect_ratio)
        image = cv2.resize(image, (new_w, new_h))
    
    # 预处理图像
    if preprocess_func:
        processed_image = preprocess_func(image)
    else:
        processed_image = image
    
    # OCR识别
    result = ocr_instance.ocr(processed_image, cls=True)
    if not result or len(result) == 0 or not result[0]:
        result = ocr_instance.ocr(image, cls=True)  # 使用原图重试
    
    # 处理OCR结果
    if result and len(result) > 0 and result[0]:
        # 按y坐标排序，确保文本从上到下
        sorted_results = sorted(result[0], key=lambda x: sum(p[1] for p in x[0])/4)
        
        boxes = []
        txts = []
        scores = []
        result_list = []
        
        for line in sorted_results:
            if float(line[1][1]) >= 0:  # 仅保留置信度较高的结果
                text = line[1][0]
                if isinstance(text, bytes):
                    text = text.decode('utf-8')
                
                # 处理坐标
                if scale_ratio:
                    # 缩放坐标回原始尺寸
                    scale_w = orig_w / new_w
                    scale_h = orig_h / new_h
                    original_box = [[int(p[0] * scale_w), int(p[1] * scale_h)] for p in line[0]]
                    # 用于绘制的坐标保持放大尺寸
                    boxes.append(line[0])
                else:
                    original_box = [[float(p[0]), float(p[1])] for p in line[0]]
                    boxes.append(original_box)
                
                txts.append(text)
                scores.append(float(line[1][1]))
                
                result_list.append({
                    "text": text,
                    "confidence": float(line[1][1]),
                    "text_region": original_box
                })
        
        return process_ocr_result(
            output_type=output_type,
            result_list=result_list,
            image=processed_image,
            boxes=boxes,
            txts=txts,
            scores=scores
        )
    else:
        return process_ocr_result(
            output_type=output_type,
            result_list=[],
            image=processed_image
        )

@app.route("/")
def home():
    """首页接口"""
    return jsonify({
        "status": "running", 
        "message": "PaddleOCR服务已启动！",
        "endpoints": {
            "ocr_system": "/predict/ocr_system（普通OCR识别）",
            "handwriting": "/predict/handwriting（手写体识别）",
            "draw_ocr": "/predict/draw_ocr（OCR可视化）"
        }
    })

@app.route('/temp/<path:filename>')
def download_file(filename):
    """文件下载接口"""
    try:
        cleanup_old_files()
        file_path = os.path.join(TEMP_DIR, filename)
        
        if not os.path.exists(file_path):
            return jsonify({
                "status": "404",
                "msg": "文件不存在"
            }), 404

        with open(file_path, 'rb') as f:
            file_content = f.read()
        os.remove(file_path)
        
        return send_file(
            BytesIO(file_content),
            mimetype='image/jpeg',
            as_attachment=False
        )
    except Exception as e:
        return jsonify({
            "status": "500",
            "msg": f"文件访问错误: {str(e)}"
        }), 500

@app.route("/predict/ocr_system", methods=["POST"])
def ocr_system():
    """普通OCR识别接口"""
    try:
        data, error = validate_request()
        if error:
            return jsonify({"status": "400", "msg": error})
        
        ocr = get_ocr(data.get("lang", "ch"))
        results = []
        
        for image_base64 in data["images"]:
            try:
                result = process_image(
                    image_base64=image_base64,
                    ocr_instance=ocr,
                    output_type=data.get("output_type", "base64"),
                    preprocess_func=preprocess_image_for_ocr
                )
                results.append(result)
            except Exception as e:
                return jsonify({"status": "500", "msg": f"图像处理错误: {str(e)}"})
        
        return jsonify({
            "status": "000",
            "msg": "success",
            "output_type": data.get("output_type", "base64"),
            "results": results
        })
    except Exception as e:
        return jsonify({"status": "500", "msg": f"服务器错误: {str(e)}"})

@app.route("/predict/draw_ocr", methods=["POST"])
def draw_ocr_api():
    """OCR可视化接口"""
    try:
        data, error = validate_request()
        if error:
            return jsonify({"status": "400", "msg": error})
        
        ocr = get_ocr(data.get("lang", "ch"))
        results = []
        
        for image_base64 in data["images"]:
            try:
                result = process_image(
                    image_base64=image_base64,
                    ocr_instance=ocr,
                    output_type=data.get("output_type", "base64"),
                    preprocess_func=preprocess_image_for_ocr
                )
                results.append(result)
            except Exception as e:
                return jsonify({"status": "500", "msg": f"图像处理错误: {str(e)}"})
        
        return jsonify({
            "status": "000",
            "msg": "success",
            "output_type": data.get("output_type", "base64"),
            "results": results
        })
    except Exception as e:
        return jsonify({"status": "500", "msg": f"服务器错误: {str(e)}"})

@app.route("/predict/handwriting", methods=["POST"])
def handwriting_recognition():
    """手写体识别接口"""
    try:
        data, error = validate_request()
        if error:
            return jsonify({"status": "400", "msg": error})
        
        ocr = get_handwriting_ocr(data.get("lang", "ch"))
        results = []
        
        for image_base64 in data["images"]:
            try:
                result = process_image(
                    image_base64=image_base64,
                    ocr_instance=ocr,
                    output_type=data.get("output_type", "base64"),
                    preprocess_func=enhance_handwriting,
                    scale_ratio=1.22  # 放大1.22倍
                )
                results.append(result)
            except Exception as e:
                return jsonify({"status": "500", "msg": f"图像处理错误: {str(e)}"})
        
        return jsonify({
            "status": "000",
            "msg": "success",
            "output_type": data.get("output_type", "base64"),
            "results": results
        })
    except Exception as e:
        return jsonify({"status": "500", "msg": f"服务器错误: {str(e)}"})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=9000, debug=False) 